#include <windows.h>
#include <commctrl.h>
#include "resource.h"
#include "Base64.h"

#define	MSGBOXTITLE		"Yonatan's Base64 Encoder/Decoder 1.0"
#define ERROR_NODLG		"Error: Could not display dialog!"
#define ERROR_NOFILEIO	"Error: Could not access file!"
#define ERROR_NOENCODE	"Error: Could not encode file!"
#define ERROR_NODECODE	"Error: Could not decode file!"
#define FILTER_SHORT	"All Files (*.*)\0*.*\0"
#define FILTER_LONG		"Yonatan's Base64 Files (*.Y64)\0*.Y64\0All Files (*.*)\0*.*\0"
#define OPEN_ENCODE		"Select File To Be Encoded..."
#define OPEN_DECODE		"Select File To Be Decoded..."
#define SAVE_ENCODE		"Save Encoded File As..."
#define SAVE_DECODE		"Save Decoded File As..."
#define Y64_DEFEXT		"Y64"

#define SPEED_ACCURACY	15
#define TIMER_INTERVAL	500
#define IDT_ENCODE		1
#define IDT_DECODE		2

HINSTANCE	hInst;
HWND		hDialog;
HANDLE		hBase64Thread = INVALID_HANDLE_VALUE;
UINT		uSpeed[SPEED_ACCURACY], uPrevFinished;
BOOL		bEncodingAsync, bCanceledAsync;
UINT		uSizeAsync;
LPBYTE		lpFileDataAsync;
LPSTR		lpBase64Async;

BOOL CALLBACK	DialogProc(HWND, UINT, WPARAM, LPARAM);
void			Encode();
void			Decode();
BOOL			LoadFile(LPBYTE*, LPDWORD, BOOL, LPCSTR, LPCSTR, LPCSTR);
BOOL			SaveFile(LPBYTE, DWORD, LPCSTR, LPCSTR, LPCSTR);
void			EnableButtons(BOOL);
void CALLBACK	Base64Callback(HANDLE, BOOL, LPVOID);
void			EncodeFile();
void			DecodeFile();
void			CancelAsyncOperation();

#ifdef _DEBUG
int WINAPI WinMain(HINSTANCE hInstance, HINSTANCE, LPSTR, int)
{
	UINT uExitCode;

	hInst = hInstance;
	uExitCode = DialogBox(hInstance, MAKEINTRESOURCE(IDD_YBASE64), NULL, DialogProc);
	CancelAsyncOperation();

	return uExitCode;
}
#else
#pragma comment(linker, "/ENTRY:Main")
void WINAPI Main()
{
	UINT uExitCode;

	hInst = GetModuleHandle(NULL);
	uExitCode = DialogBox(hInst, MAKEINTRESOURCE(IDD_YBASE64), NULL, DialogProc);
	CancelAsyncOperation();

	ExitProcess(uExitCode);
}
#endif // !_DEBUG

BOOL CALLBACK DialogProc(HWND hDlg, UINT uMsg, WPARAM wParam, LPARAM)
{
	UINT uFinished, uTotal, i, uAverageSpeed, uNonzeroCount, uPercent;
	CHAR lpMessage[100];
	BOOL bSuccess;
	
	switch(uMsg) {
	case WM_CLOSE:
		EndDialog(hDlg, 0);
		break;
	
	case WM_INITDIALOG:
		SendMessage(hDialog = hDlg, WM_SETICON, ICON_BIG,
			(LPARAM)LoadIcon(hInst, MAKEINTRESOURCE(IDI_MAINICON)));
		SendDlgItemMessage(hDlg, IDC_ASYNC, BM_SETCHECK, BST_CHECKED, NULL);
		SendDlgItemMessage(hDlg, IDC_PROCESS, PBM_SETRANGE, 0, 100);

		break;
		
	case WM_COMMAND:
		if(wParam == MAKEWPARAM(IDC_PLAINTEXT, EN_CHANGE))
			Encode();

		if(wParam == MAKEWPARAM(IDC_BASE64TEXT, EN_CHANGE))
			Decode();

		if(wParam == MAKEWPARAM(IDC_ENCODEFILE, BN_CLICKED))
			EncodeFile();

		if(wParam == MAKEWPARAM(IDC_DECODEFILE, BN_CLICKED))
			DecodeFile();

		if(wParam == MAKEWPARAM(IDC_CANCEL, BN_CLICKED))
			CancelAsyncOperation();

		break;

	case WM_TIMER:
		if(wParam == IDT_ENCODE)
			bSuccess = GetEncodeAsyncState(hBase64Thread, &uFinished, &uTotal);

		if(wParam == IDT_DECODE)
			bSuccess = GetDecodeAsyncState(hBase64Thread, &uFinished, &uTotal);

		if(!bSuccess)
			break;

		for(i = 0; i < SPEED_ACCURACY - 1; i++)
			uSpeed[i] = uSpeed[i + 1];

		uSpeed[SPEED_ACCURACY - 1] = uFinished - uPrevFinished;
		uPrevFinished = uFinished;

		for(i = uAverageSpeed = uNonzeroCount = 0; i < SPEED_ACCURACY; uAverageSpeed += uSpeed[++i])
			if(uSpeed[i])
				uNonzeroCount++;
		
		uAverageSpeed /= (uNonzeroCount << 10);
		uFinished >>= 10;
		uTotal >>= 10;

		if(uTotal)
			uPercent = 100 * uFinished / uTotal;
		else
			uPercent = 100;
		
		wsprintf(lpMessage, "%d KB %scoded out of %d KB total (%d%%), average speed: %d KB/s.",
			uFinished, wParam == IDT_ENCODE ? "en" : "de", uTotal, uPercent, uAverageSpeed);
		
		SendDlgItemMessage(hDlg, IDC_PROCESS, PBM_SETPOS, uPercent, NULL);
		SetDlgItemText(hDlg, IDC_PERCENT, lpMessage);

		break;
		
	default:
		return FALSE;
	}

	return TRUE;
}

void Encode()
{
	LPSTR	lpText, lpBase64;
	UINT	uLenText, uLenBase64;

	uLenText = SendDlgItemMessage(hDialog, IDC_PLAINTEXT, WM_GETTEXTLENGTH, NULL, NULL);
	lpText = (LPSTR)LocalAlloc(LPTR, uLenText + 1);
	GetDlgItemText(hDialog, IDC_PLAINTEXT, lpText, uLenText + 1);
	
	uLenBase64 = CharsNeededToEncode(uLenText);
	lpBase64 = (LPSTR)LocalAlloc(LPTR, uLenBase64);

	Base64Encode((LPCBYTE)lpText, uLenText, lpBase64, uLenBase64);

	LocalFree(lpText);
	SetDlgItemText(hDialog, IDC_BASE64TEXT, lpBase64);
	LocalFree(lpBase64);
}

void Decode()
{
	LPSTR	lpBase64, lpText;
	UINT	uLenBase64, uEqualSigns, uLenText;
	
	uLenBase64 = SendDlgItemMessage(hDialog, IDC_BASE64TEXT, WM_GETTEXTLENGTH, NULL, NULL);
	lpBase64 = (LPSTR)LocalAlloc(LPTR, uLenBase64 + 1);
	GetDlgItemText(hDialog, IDC_BASE64TEXT, lpBase64, uLenBase64 + 1);

	uEqualSigns = CountEqualSigns(lpBase64);
	uLenText = BytesNeededToDecode(lstrlen(lpBase64), uEqualSigns);
	lpText = (LPSTR)LocalAlloc(LPTR, uLenText + 1);

	if(!Base64Decode(lpBase64, (LPBYTE)lpText, uLenText)) {
		SetDlgItemText(hDialog, IDC_PLAINTEXT, NULL);
		LocalFree(lpBase64);
		return;
	}
	
	LocalFree(lpBase64);
	SetDlgItemText(hDialog, IDC_PLAINTEXT, lpText);
	LocalFree(lpText);
}

BOOL LoadFile(LPBYTE* lpFileData, LPDWORD lpFileSize, BOOL bAllocateExtraByte,
			  LPCSTR lpFilter, LPCSTR lpTitle, LPCSTR lpDefExt)
{
	OPENFILENAME ofn = {
		sizeof(OPENFILENAME),									// lStructSize
		hDialog,												// hWndOwner
		hInst,													// hInstance
		lpFilter,												// lpstrFilter
		NULL,													// lpstrCustomFilter
		NULL,													// nMaxCustFilter
		1,														// nFilterIndex
		(LPSTR)LocalAlloc(LPTR, MAX_PATH),						// lpstrFile
		MAX_PATH,												// nMaxFile
		NULL,													// lpstrFileTitle
		NULL,													// nMaxFileTitle
		NULL,													// lpstrInitialDir
		lpTitle,												// lpstrTitle
		OFN_FILEMUSTEXIST | OFN_HIDEREADONLY | OFN_LONGNAMES,	// Flags
		NULL,													// nFileOffset
		NULL,													// nFileExtension
		lpDefExt,												// lpstrDefExt
		NULL,													// lCustData
		NULL,													// lpfnHook
		NULL													// lpTemplateName
	};

	HANDLE	hFile;
	DWORD	dwBytesRead;
	BOOL	bFailure;

	if(!GetOpenFileName(&ofn)) {
		if(CommDlgExtendedError())
			MessageBox(hDialog, ERROR_NODLG, MSGBOXTITLE, MB_ICONSTOP);
		LocalFree(ofn.lpstrFile);
		return FALSE;
	}

	bFailure = ((hFile = CreateFile(ofn.lpstrFile, GENERIC_READ, FILE_SHARE_READ,
		NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)) == INVALID_HANDLE_VALUE);
	LocalFree(ofn.lpstrFile);
	
	if(bFailure) {
		MessageBox(hDialog, ERROR_NOFILEIO, MSGBOXTITLE, MB_ICONSTOP);
		return FALSE;
	}

	if((*lpFileSize = GetFileSize(hFile, NULL)) == INVALID_FILE_SIZE) {
		CloseHandle(hFile);
		MessageBox(hDialog, ERROR_NOFILEIO, MSGBOXTITLE, MB_ICONSTOP);
		return FALSE;
	}
	
	*lpFileData = (LPBYTE)LocalAlloc(LPTR, *lpFileSize + bAllocateExtraByte);
	bFailure = !ReadFile(hFile, *lpFileData, *lpFileSize, &dwBytesRead, NULL)
		|| *lpFileSize != dwBytesRead;
	CloseHandle(hFile);

	if(bFailure) {
		LocalFree(*lpFileData);
		MessageBox(hDialog, ERROR_NOFILEIO, MSGBOXTITLE, MB_ICONSTOP);
		return FALSE;
	}

	return TRUE;
}

BOOL SaveFile(LPBYTE lpFileData, DWORD dwFileSize, LPCSTR lpFilter, LPCSTR lpTitle, LPCSTR lpDefExt)
{
	OPENFILENAME ofn = {
		sizeof(OPENFILENAME),									// lStructSize
		hDialog,												// hWndOwner
		hInst,													// hInstance
		lpFilter,												// lpstrFilter
		NULL,													// lpstrCustomFilter
		NULL,													// nMaxCustFilter
		1,														// nFilterIndex
		(LPSTR)LocalAlloc(LPTR, MAX_PATH),						// lpstrFile
		MAX_PATH,												// nMaxFile
		NULL,													// lpstrFileTitle
		NULL,													// nMaxFileTitle
		NULL,													// lpstrInitialDir
		lpTitle,												// lpstrTitle
		OFN_HIDEREADONLY | OFN_LONGNAMES | OFN_OVERWRITEPROMPT,	// Flags
		NULL,													// nFileOffset
		NULL,													// nFileExtension
		lpDefExt,												// lpstrDefExt
		NULL,													// lCustData
		NULL,													// lpfnHook
		NULL													// lpTemplateName
	};

	HANDLE	hFile;
	DWORD	dwBytesWritten;
	BOOL	bFailure;

	if(!GetSaveFileName(&ofn)) {
		if(CommDlgExtendedError())
			MessageBox(hDialog, ERROR_NODLG, MSGBOXTITLE, MB_ICONSTOP);
		LocalFree(ofn.lpstrFile);
		return FALSE;
	}

	bFailure = ((hFile = CreateFile(ofn.lpstrFile, GENERIC_WRITE, FILE_SHARE_READ,
		NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL)) == INVALID_HANDLE_VALUE);
	LocalFree(ofn.lpstrFile);

	if(bFailure) {
		MessageBox(hDialog, ERROR_NOFILEIO, MSGBOXTITLE, MB_ICONSTOP);
		return FALSE;
	}

	bFailure = !WriteFile(hFile, lpFileData, dwFileSize, &dwBytesWritten, NULL)
		|| dwFileSize != dwBytesWritten;
	CloseHandle(hFile);
	
	if(bFailure) {
		MessageBox(hDialog, ERROR_NOFILEIO, MSGBOXTITLE, MB_ICONSTOP);
		return FALSE;
	}

	return TRUE;
}

void EnableButtons(BOOL bEnable)
{
	EnableWindow(GetDlgItem(hDialog, IDC_ENCODEFILE), bEnable);
	EnableWindow(GetDlgItem(hDialog, IDC_DECODEFILE), bEnable);
	EnableWindow(GetDlgItem(hDialog, IDC_ASYNC), bEnable);
	EnableWindow(GetDlgItem(hDialog, IDC_CANCEL), !bEnable);
}

void CALLBACK Base64Callback(HANDLE hThread, BOOL bSuccess, LPVOID)
{
	if(bEncodingAsync) {
		KillTimer(hDialog, IDT_ENCODE);

		if(bSuccess)
			SaveFile((LPBYTE)lpBase64Async, uSizeAsync - 1, FILTER_LONG, SAVE_ENCODE, Y64_DEFEXT);
		else if(!bCanceledAsync)
			MessageBox(hDialog, ERROR_NOENCODE, MSGBOXTITLE, MB_ICONSTOP);
	} else {
		KillTimer(hDialog, IDT_DECODE);

		if(bSuccess)
			SaveFile(lpFileDataAsync, uSizeAsync, FILTER_SHORT, SAVE_DECODE, NULL);
		else if(!bCanceledAsync)
			MessageBox(hDialog, ERROR_NODECODE, MSGBOXTITLE, MB_ICONSTOP);
	}

	EnableButtons(TRUE);
	CloseHandle(hThread);
	hBase64Thread = INVALID_HANDLE_VALUE;
	LocalFree(lpFileDataAsync);
	LocalFree(lpBase64Async);
	SendDlgItemMessage(hDialog, IDC_PROCESS, PBM_SETPOS, 0, NULL);
	SetDlgItemText(hDialog, IDC_PERCENT, NULL);
	SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_ARROW)));
}

void EncodeFile()
{
	LPBYTE	lpFileData;
	DWORD	dwFileSize;
	LPSTR	lpBase64;
	UINT	uLenBase64, i;
	BOOL	bFailure, bAsync;

	if(!LoadFile(&lpFileData, &dwFileSize, FALSE, FILTER_SHORT, OPEN_ENCODE, NULL))
		return;

	uLenBase64 = CharsNeededToEncode(dwFileSize);
	lpBase64 = (LPSTR)LocalAlloc(LPTR, uLenBase64);
	bAsync = (SendDlgItemMessage(hDialog, IDC_ASYNC, BM_GETCHECK, NULL, NULL) == BST_CHECKED);
	SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_WAIT)));

	if(bAsync) {
		EnableButtons(FALSE);
		bEncodingAsync = TRUE;
		bCanceledAsync = FALSE;
		lpFileDataAsync = lpFileData;
		lpBase64Async = lpBase64;
		uSizeAsync = uLenBase64;
		uPrevFinished = 0;
		for(i = 0; i < SPEED_ACCURACY; i++)
			uSpeed[i] = 0;
		SetTimer(hDialog, IDT_ENCODE, TIMER_INTERVAL, NULL);
		
		bFailure = ((hBase64Thread = Base64EncodeAsync(lpFileData, dwFileSize, lpBase64,
			uLenBase64, Base64Callback, NULL)) == INVALID_HANDLE_VALUE);
	} else {
		bFailure = !Base64Encode(lpFileData, dwFileSize, lpBase64, uLenBase64);
		SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_ARROW)));
	}

	if(bFailure) {
		LocalFree(lpFileData);
		LocalFree(lpBase64);
		SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_ARROW)));
		MessageBox(hDialog, ERROR_NOENCODE, MSGBOXTITLE, MB_ICONSTOP);
	}
	
	if(bAsync || bFailure)
		return;

	SaveFile((LPBYTE)lpBase64, uLenBase64 - 1, FILTER_LONG, SAVE_ENCODE, Y64_DEFEXT);
	LocalFree(lpBase64);
	LocalFree(lpFileData);
}

void DecodeFile()
{
	LPSTR	lpBase64;
	UINT	uLenBase64, uEqualSigns, i;
	LPBYTE	lpFileData;
	DWORD	dwFileSize;
	BOOL	bFailure, bAsync;

	if(!LoadFile((LPBYTE*)&lpBase64, (LPDWORD)&uLenBase64, TRUE, FILTER_LONG, OPEN_DECODE, Y64_DEFEXT))
		return;

	uEqualSigns = CountEqualSigns(lpBase64);
	dwFileSize = BytesNeededToDecode(uLenBase64, uEqualSigns);
	lpFileData = (LPBYTE)LocalAlloc(LPTR, dwFileSize);
	bAsync = (SendDlgItemMessage(hDialog, IDC_ASYNC, BM_GETCHECK, NULL, NULL) == BST_CHECKED);
	SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_WAIT)));

	if(bAsync) {
		EnableButtons(FALSE);
		bEncodingAsync = FALSE;
		bCanceledAsync = FALSE;
		lpFileDataAsync = lpFileData;
		lpBase64Async = lpBase64;
		uSizeAsync = dwFileSize;
		uPrevFinished = 0;
		for(i = 0; i < SPEED_ACCURACY; i++)
			uSpeed[i] = 0;
		SetTimer(hDialog, IDT_DECODE, TIMER_INTERVAL, NULL);
		
		bFailure = ((hBase64Thread = Base64DecodeAsync(lpBase64, lpFileData,
			dwFileSize, Base64Callback, NULL)) == INVALID_HANDLE_VALUE);
	} else {
		bFailure = !Base64Decode(lpBase64, lpFileData, dwFileSize);
		SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_ARROW)));
	}

	if(bFailure) {
		LocalFree(lpFileData);
		LocalFree(lpBase64);
		SetCursor(LoadCursor(NULL, MAKEINTRESOURCE(IDC_ARROW)));
		MessageBox(hDialog, ERROR_NODECODE, MSGBOXTITLE, MB_ICONSTOP);
	}

	if(bAsync || bFailure)
		return;

	SaveFile(lpFileData, dwFileSize, FILTER_SHORT, SAVE_DECODE, NULL);
	LocalFree(lpFileData);
	LocalFree(lpBase64);
}

void CancelAsyncOperation()
{
	if(hBase64Thread == INVALID_HANDLE_VALUE)
		return;

	TerminateThread(hBase64Thread, 0);
	bCanceledAsync = TRUE;
	Base64Callback(hBase64Thread, FALSE, NULL);
}